package ch.epfl.lara.pm.transformation;

import ch.epfl.lara.pm.ast._
import ch.epfl.lara.pm.util.{MyImplicitConversions,Substitution}
import ch.epfl.lara.pm.analyze.Util

import scala.collection.mutable.{HashMap,Map}

object MatchAnalyzer extends MyImplicitConversions {
  
  type Path = String
  type Variable = String
  type Substitution = (Expr,Variable)
  
  def apply(tree: Tree): Tree = tree match {
    case Program(constr,objs,clss,main) => 
      Program(constr,objs.map(this(_).asInstanceOf[ObjectDef]),clss.map(this(_).asInstanceOf[ClassDef]),this(main).asInstanceOf[Expr])  
    
    case ObjectDef(name,extr) => 
      ObjectDef(name,this(extr).asInstanceOf[ExtractorDef])
      
    case ExtractorDef(arg,tpe,body) => 
      ExtractorDef(arg,tpe,this(body).asInstanceOf[Expr])
      
    case ClassDef(isAbs,name,fields,superc,methods) => 
      ClassDef(isAbs,name,fields,superc,methods.map(this(_).asInstanceOf[MethodDef]))
      
    case pd @ ParamDef(name,tpe) => pd
    
    case MethodDef(name,params,tpe,body) => 
      MethodDef(name,params,tpe,this(body).asInstanceOf[Expr])
      
    case m: Match => 
      analyzeMatch(m,HashMap.empty,HashMap.empty).asInstanceOf[Match]
    
    case _ => tree
  }
  
  private def analyzeMatch(m: PatternMatching, pathEnv: HashMap[Path,Variable], guardEnv: HashMap[Expr,Var]): PatternMatching = m match {
    case Match(scrutinee,tpe,cases) => 
      Match(scrutinee,tpe,cases.map(analyzeMatch(_,pathEnv,guardEnv).asInstanceOf[CaseDef]))
    
    case cd @ CaseDef(_,pattern,guard,expr) =>
      analyzeCase(cd,pathEnv,guardEnv)
      
    case _ => m
  }
  
  private def analyzeCase(kase: CaseDef, pathEnv: HashMap[Path,Variable], guardEnv: HashMap[Expr,Var]): CaseDef = {
    val CaseDef(pos,pattern,guard,expr) = kase
    analyzePatternPath(pattern)
    val patternSubst = collectVariablePatternSubstitutions(pattern,pathEnv)
    val nguard = (if(guard.isDefined) Some(analyzeGuard(guard.get,guardEnv))
                 else guard)
    Substitution(CaseDef(pos,pattern,nguard,expr),patternSubst)
  }
  
  
  private def analyzeGuard(guard: Expr, guardEnv: HashMap[Expr,Var]): Expr = guard match {
    case True | False | IsInstanceOf(_,_) | Number(_) | Var(_) | INone => guard
    case Call(obj,method,args) => 
      guardEnv.get(guard) match {
        case None => 
          val fresh = Var(Util.freshVarName(method))
          guardEnv.update(guard,fresh)
          fresh
        
        case Some(variable) => variable
      }
      
    case Select(obj,field) => 
      guardEnv.get(guard) match {
        case None => 
          val fresh = Var(Util.freshVarName(field))
          guardEnv.update(guard,fresh)
          fresh
    
        case Some(variable) => variable
      }
      
    case New(name,args) => New(name,args.map(analyzeGuard(_,guardEnv)))
    
    case Plus(left,right) => Plus(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Minus(left,right) => Minus(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Prod(left,right) => Prod(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Div(left,right) => Div(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Mod(left,right) => Mod(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Eq(left,right) => Eq(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Less(left,right) => Less(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case LessEq(left,right) => LessEq(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case EAnd(left,right) => EAnd(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case EOr(left,right) => EOr(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case Not(expr) => Not(analyzeGuard(expr,guardEnv))
    case NotEq(left,right) => NotEq(analyzeGuard(left,guardEnv),analyzeGuard(right,guardEnv))
    case _ => guard
  }
  
  private def analyzePatternPath(pattern: PatternDef): Unit = { 
    // store the path to the analyzed pattern 
    def analyzePath(pattern: PatternDef, path: String): Unit = pattern match {
      case v @ VarPattern(name,tpe) => 
        name.path = path
      case ClassPattern(objname,fields) => 
        objname.path = path
        fields.zipWithIndex.map(pair => analyzePath(pair._1,path+objname+pair._2))
    }
    analyzePath(pattern,"")
  }
  
  
  
  private def collectVariablePatternSubstitutions(pattern: PatternDef, env: HashMap[Path,Variable]): List[Substitution] = {
    def substitution(varname: PatternName): Substitution = env.get(varname.path) match {
      case None => 
        val freshVarname = Util.freshVarName(varname)
        env += ((varname.path,freshVarname))
        (Var(varname),freshVarname)
    
     case Some(substVariable) => (Var(varname),substVariable)
    }
    
    def collect(pattern: PatternDef): List[Substitution] = pattern match {
      case VarPattern(name,tpe) => List(substitution(name)) 
      case ClassPattern(name,fields) => fields.flatMap(collect(_))
    }
    
    collect(pattern)
  }

 
}
 